#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <signal.h>
#include <netinet/in.h>
#include <netdb.h> 
#include <time.h>

#define TRUE 1
#define FALSE 0
#define FAILED -1
#define BAD_ID 999999
//#define _CLIENT_DEBUG_

const char PATTERN = 0x5A; //01011010 = Z

struct my_timer
{
   int valid,retransmissions;
   unsigned int message_id;
   struct timeval initial_time;
   struct timeval timer;
};

extern int    errno;
extern int    optind;
extern char * optarg;

int neg_ack = FALSE;
int windowing = FALSE;
int window = 1;
int retransmission_time = 0;
int number_of_packets;
int size_of_packets;
int alarm_pipe_fd[2];

struct my_timer *timers;

char Usage[] = "\
Usage: traffic_generator [-options] to_host to_port n_packets n_size\n\
Options:\n\
      -r ## Retransmission timer for sent packets in microseconds\n\
      -w ## The number of outstanding unacknowledged packets allowed\n\
            also known as the window size\n\
      -n    Signals the program to look for negative acknowledgements\n\
            when a message has been changed during transfer\n";

void leave(const char *msg)
{
    puts(msg);
    exit(1);
}

void error(const char *msg)
{
    perror(msg);
    exit(2);
}

void usage(const char *msg)
{
   puts(msg);
   puts(Usage);
   exit(3);
}

long long timeval_diff(struct timeval *difference,
                       struct timeval *start_time,
                       struct timeval *end_time)
{
   struct timeval temp_diff;

   if(difference==NULL)
   {
      difference=&temp_diff;
   }

   difference->tv_sec =end_time->tv_sec -start_time->tv_sec ;
   difference->tv_usec=end_time->tv_usec-start_time->tv_usec;

   // Using while instead of if below makes the code slightly more robust. 

   while(difference->tv_usec<0)
   {
      difference->tv_usec+=1000000;
      difference->tv_sec -=1;
   }

   return 1000000LL*difference->tv_sec+
                    difference->tv_usec;
}

//given a timer index, set the timer for when a retransmission needs to occur.
void set_transmission_timer(int timer_num)
{   
   if (gettimeofday(&timers[timer_num].timer,NULL))
      error("ERROR, getting time");   
   
   timers[timer_num].timer.tv_usec += retransmission_time;
   
   while (timers[timer_num].timer.tv_usec > 1000000)
   {
      timers[timer_num].timer.tv_sec += 1;
      timers[timer_num].timer.tv_usec -= 1000000;
   }
}

//used to find an index of a timer in the timers list given a message id 
int get_timer_by_id(int message_id)
{
   int i;
   for (i = 0; i < window; i++)
      if (timers[i].valid && timers[i].message_id == message_id)
         return i;
   return -1;
}
   
int main(int argc, char* argv[])
{
   const int SIZE_OF_RECV_DATA = 100;
   const int SIZE_OF_RESULTS = 1000;
   const int TEMP_RESULTS_LENGTH = 100;   

   fd_set fds,master;
   int fdmax;
   
   long to_port;
   struct hostent *to;   
   int sockfd;   
   
   long long running_sum = 0;  
   long long packet_time = 0;
   long long max = 0,min = 100000000;   
   int total_retransmissions = 0,nak_retransmissions = 0, timeout_retransmissions = 0;
   int bit_rate,pps;
   unsigned int id,message_id_offset = 0;
   unsigned int *id_pt,*finished_ids,*resend_ids;
   int n,nbytes,i,itemp,j,jtemp,num,sent = 0,received = 0,done = FALSE,send_counter = 0;
   char *ctemp;
   char *packet_data,recv_data[SIZE_OF_RECV_DATA];
   char results[SIZE_OF_RESULTS],temp_results[100];
   char c;
   struct sockaddr_in to_addr;

   struct timeval tstart,tfinish,tdiff,tnow,twait;
   struct itimerval time_out;
   
   //assign initial values to stuff
   FD_ZERO(&fds);
   bzero((char *) &to_addr, sizeof(to_addr));
   
   // check for the correct inputs
   while((c = getopt(argc, argv, "w:r:n")) != -1)
   {
      switch (c)
      {
         case 'w':
            windowing = TRUE;
            window = strtol(optarg,&ctemp,0);
            if (*ctemp != NULL || window < 1 || window > 100)
               usage("Invalid Window Size, please select 1-100.");            
            break;
         case 'n':
            neg_ack = TRUE;
            break;
         case 'r':
            retransmission_time = strtol(optarg,&ctemp,0);
            if (*ctemp != NULL || retransmission_time < 1 || retransmission_time > 10000000)
               usage("Invalid retransmission time! Select between 1 and 10,000,000 us");
            break;
         default:
            usage("Invalid options chosen");
      }
   }
   
   if (argc - optind < 4)
   {
      usage("Invalid amount of arguments!");
   }
   else
   {
      to = gethostbyname(argv[optind]);
      if (to == NULL)
         error("ERROR, no such host");
      
      to_port = strtol(argv[optind+1],&ctemp,0);
      if (*ctemp != NULL || to_port > 65535 || to_port < 1024)
         leave("Invalid port number");
      
      number_of_packets = strtol(argv[optind+2],&ctemp,0);
      if (*ctemp != NULL || number_of_packets < 1 || number_of_packets > 10000000) //limit saves program from running forever
         leave("Invalid number of packets - must be 1 <= packets <= 10000000 ");
      
      size_of_packets = strtol(argv[optind+3],&ctemp,0);
      if (*ctemp != NULL || size_of_packets < 5 || size_of_packets > 50000)
         leave("Invalid size of packets - must be between 5 and 50000");   
   }      

//just using this to print out what the arguments are so that I can test
#ifdef _CLIENT_DEBUG_  
   fprintf(stderr,"Using %s as the host name.\n", argv[optind]);
   fprintf(stderr,"Using port %d.\n", to_port);
   fprintf(stderr,"Using packet size: %d\n", size_of_packets);
   fprintf(stderr,"Using %d iterations per packet size.\n", number_of_packets);
   fprintf(stderr,"Window Size: %d\n", window);
   if (neg_ack)
      fprintf(stderr,"Using negative acknowledgements.\n");
  
   exit(1);
#endif   
   
   //set up timers
   timers = malloc(sizeof(struct my_timer) * window);  
   finished_ids = malloc(sizeof(unsigned int) * window);
   resend_ids = malloc(sizeof(unsigned int) * window);
   for (i = 0; i < window; i++)
   {
      timers[i].valid = FALSE;   
      finished_ids[i] = BAD_ID;
      resend_ids[i] = BAD_ID;
   }
   
   //set up socket
   if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) 
      error("ERROR opening socket");    
   
   //connect socket
   to_addr.sin_family = AF_INET;
   bcopy((char *) to->h_addr, (char *) &to_addr.sin_addr.s_addr, to->h_length);
   to_addr.sin_port = htons(to_port);
   
   if (connect(sockfd,(struct sockaddr*) &to_addr, sizeof(to_addr)) < 0)
      error("ERROR connecting");
      
   packet_data = malloc(size_of_packets);
   bzero(packet_data,size_of_packets);
   memset(packet_data,PATTERN,size_of_packets); 
   
   //set up our fd_set
   FD_SET(sockfd, &master);
   
   if (sockfd > alarm_pipe_fd[1])
      fdmax = sockfd + 1;
   else 
      fdmax = alarm_pipe_fd[1] + 1;
   
   //setting our id pointer to point to the place in the recv_data buffer that the id will be
   id_pt = (int*)packet_data;
   
   //set up timer
   twait.tv_sec = 0;
   twait.tv_usec = 0;
   
   //this counter is useful for the quick packet turn around tiems discussed in the sending data
   //section just below.
   send_counter = 0;
   
   while (!done)
   {
      fds = master;

      //fill up the pipeline with messages at every chance 
      
      //debatable how to do this - if there's a longer wait time like
      //sending data across the internet and the delay from sending data
      //across the network is greater than that of our own processing
      //you should use the for loop and have many packets in flight.  
      
      //other method is prioritizing the receiving of messages and spending
      //less time sending.  This works better for local networks with low 
      //packet error and drop rates because it is more rapidly ready to receive
      //data from the socket, though this method may be a bit slower for 
      //slower networks - it won't be much because there won't be any data
      //to be received for a while so you will loop through many times 
      //covering all the pieces relatively quickly.  This is similar to a 
      //clock scheduler that you check certain things at each time around.
      
      //for (i = 0; i < window; i++)
      i = send_counter;
      {
         //if the timer is valid and we are doing retransmissions check 
         //this guys retransmission timer to see if we need to send again.
         if (timers[i].valid && retransmission_time > 0)
         {
            if (gettimeofday(&tnow,NULL))
               error("ERROR, get start time");     
            
            if (timeval_diff(&tdiff,&tnow,&timers[i].timer) < 0)
            {   
               timers[i].retransmissions++;
               timeout_retransmissions++;
               set_transmission_timer(i);
               
               id_pt[0] = htonl(timers[i].message_id);
               
               write(sockfd, packet_data, size_of_packets);
               // packet_data[size_of_packets-1] = '\0';
               // ctemp = packet_data + 4; 
               // puts(ctemp);
            }
         }
         //if the timer isn't valid that means we have a timer that can be used
         //to send a new message.
         else if (!timers[i].valid && sent != number_of_packets)
         {
            if (gettimeofday(&tnow,NULL))
               error("ERROR, get start time");         
            
            //reset the data for this one
            timers[i].valid = TRUE;
            timers[i].retransmissions = 0;
            timers[i].message_id = sent + message_id_offset;
            
            //set up timing data
            timers[i].initial_time.tv_sec  = tnow.tv_sec;
            timers[i].initial_time.tv_usec = tnow.tv_usec; 
            set_transmission_timer(i);
            
            //send data
            id_pt[0] = htonl(timers[i].message_id);
            write(sockfd, packet_data, size_of_packets);  
            
            sent++;
         }
      }
      
      //grab the current time that all the following packets arrived at (or close enough)
      if (gettimeofday(&tnow,NULL))
         error("ERROR, get finish time");          
      
      //while there are packets to be received get all of them and add their message_id
      //to the received_ids list so that we can establish 
      itemp = jtemp = 0;      
      while (select(fdmax,&fds,NULL,NULL,&twait) > 0) //means the only other option is that we received a message on our socket
      {         
         if ((nbytes = read(sockfd, recv_data, SIZE_OF_RECV_DATA)) < 0)
            error("ERROR, reading from the socket");
         
         //puts("checking ack");
         if (neg_ack)
         {   
            if (strncmp(recv_data + 4, "NAK",3) == 0)
               resend_ids[jtemp++] = ntohl(*(unsigned int*)recv_data); 
            else
               finished_ids[itemp++] = ntohl(*(unsigned int*)recv_data); 
         }
         else
         {
            finished_ids[itemp++] = ntohl(*(unsigned int*)recv_data);            
         }
      }
      
      //loop through and resends any nak'd packets
      //puts("Checking need to resend");
      n = 0;
      while (n < jtemp)
      {
         //try to get the timer associated with the message_id received if 
         //its not valid just leave
         if ((num = get_timer_by_id(resend_ids[n++])) == -1)
            continue;
         
         timers[num].retransmissions++;
         nak_retransmissions++;
         set_transmission_timer(num);
         
         id_pt[0] = htonl(timers[num].message_id);
         
         write(sockfd, packet_data, size_of_packets);   
      }
      
      //loop through and handle all received messages
      //puts("Grabbing valid data");
      n = 0;
      while (n < itemp)
      {
         //check if the received message is for a currently valid timer
         //if not then just continue searching.
         if ((num = get_timer_by_id(finished_ids[n++])) == -1)      
            continue;               
         
         //invalidate timer so it can be reused for other packets
         //printf("Invalidating timer - %d for message %u\n",num,id);
         received++;
         //printf("Received Message - %d\n",timers[num].message_id);
         timers[num].valid = FALSE;
         total_retransmissions += timers[num].retransmissions;
         
         //check timing data
         //puts("Performing Timing Calculations");
         packet_time = timeval_diff(&tdiff,&timers[num].initial_time,&tnow);
         running_sum += packet_time;
         
         //printf("Took %lld us to send packet\n", packet_time);
         if (packet_time > max)
            max = packet_time;
      
         if (packet_time < min)
            min = packet_time;  
      }
      
      //check if we have gotten the amount of packets expected to
      if (received >= number_of_packets)
         break;     
      
      //increment our send counter for tracking of the hand of the clock.
      send_counter = (send_counter + 1) % window;
   }    

   //calculate statistics   
   bit_rate = 8 * size_of_packets / ((running_sum/number_of_packets) / 1E6); //kinda?
   pps = (float)window * (1E6f/((float)running_sum/(float)number_of_packets));
   
   //generate results text
   results[0] = '\0';
  
   snprintf(temp_results,TEMP_RESULTS_LENGTH,"Summary:\nNumber of %d byte packets sent and acknowledged: %d\n", size_of_packets, number_of_packets);
   strcat(results,temp_results); 
   if (windowing)
   {
      snprintf(temp_results,TEMP_RESULTS_LENGTH,"Window size: %d\n",window);
      strcat(results,temp_results);
   }
   
   snprintf(temp_results,TEMP_RESULTS_LENGTH,"Average time to send packet: %lld us\n", running_sum/number_of_packets);
   strcat(results,temp_results);
   snprintf(temp_results,TEMP_RESULTS_LENGTH,"Max time to send packet: %ld us\n", max);
   strcat(results,temp_results);
   snprintf(temp_results,TEMP_RESULTS_LENGTH,"Min time to send packet: %ld us\n", min);
   strcat(results,temp_results);
   snprintf(temp_results,TEMP_RESULTS_LENGTH,"Average packets per second: %d pps\n", pps);
   strcat(results,temp_results);
   if (retransmission_time > 0 || neg_ack)
   {
      snprintf(temp_results,TEMP_RESULTS_LENGTH,"Retransmissions: %d (%d timeout, %d nak)\n",total_retransmissions,timeout_retransmissions,nak_retransmissions);
      strcat(results,temp_results);
   }
   
   //tell the other side your results
   write(sockfd, results, strlen(results));
   puts(results);
   
   //clean up - for some reason it is cratering on some of these
   //I have no clue as to why so now I've just resorted to not cleaning up.
   //free(packet_data);
   //free(timers);
   //free(finished_ids);
   //free(resend_ids);
   
   return 0;
}
